import warnings
from argparse import ArgumentParser
import os

import mmcv
import models  # noqa: F401,F403
from mmcls.apis import inference_model, show_result_pyplot
from mmcls.models import build_classifier
from mmcv.runner import load_checkpoint


def init_model(config, checkpoint=None, device='cuda:0'):
    """Initialize a classifier from config file.

    Args:
        config (str or :obj:`mmcv.Config`): Config file path or the config
            object.
        checkpoint (str, optional): Checkpoint path. If left as None, the model
            will not load any weights.

    Returns:
        nn.Module: The constructed classifier.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {type(config)}')
    # config.model.pretrained = None
    model = build_classifier(config.model)
    # if checkpoint is not None:
    #     map_loc = 'cpu' if device == 'cpu' else None
    #     checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc)
    #     if 'CLASSES' in checkpoint['meta']:
    #         model.CLASSES = checkpoint['meta']['CLASSES']
    #     else:
    #         from mmcls.datasets import ImageNet
    #         warnings.simplefilter('once')
    #         warnings.warn('Class names are not saved in the checkpoint\'s '
    #                       'meta data, use imagenet by default.')
    #         model.CLASSES = ImageNet.CLASSES

    from mmcls.datasets import ImageNet
    warnings.simplefilter('once')
    warnings.warn('Class names are not saved in the checkpoint\'s '
                  'meta data, use imagenet by default.')
    model.CLASSES = ImageNet.CLASSES

    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()
    return model


def main():
    parser = ArgumentParser()
    parser.add_argument("img", help="Image file")
    parser.add_argument("config", help="Config file")
    parser.add_argument("checkpoint", help="Checkpoint file")
    parser.add_argument("--device",
                        default="cuda:0",
                        help="Device used for inference")
    args = parser.parse_args()

    checkpoint = args.checkpoint
    if isinstance(checkpoint, str) and os.path.exists(checkpoint):
        # build the model from a config file and a checkpoint file
        model = init_model(args.config, args.checkpoint, device=args.device)
    else:
        model = init_model(args.config, None, device=args.device)
    # test a single image
    result = inference_model(model, args.img)
    print(result)
    # show the results
    show_result_pyplot(model, args.img, result)


if __name__ == "__main__":
    main()
